set.seed(02138)

pi0 <- 0.4
p <- 0.2
r <- 20
n <- 100000

ind <- sample(c(0, 1), n, prob = c(pi0, 1 - pi0), replace = TRUE)
Z <- ind*rnbinom(n, mu = p*r/(1-p), size = r) 
S <- sd(Z)
X <- Z + rnorm(n, 0, sd = S)


moments_df <- Rmoments(X, S, R = r, n = n)

norm <- paramsNormal(X, S, moments_df = moments_df, plot_dp = FALSE)
pois <- paramsPoisson(X, S, moments_df = moments_df, plot_dp = FALSE)
nbinom <- paramsNB(X, S, moments_df = moments_df, plot_dp = FALSE)
zip <- paramsZIP(X, S, moments_df = moments_df, plot_dp = FALSE)


x <- paramsZINB(X, S, moments_df = moments_df, plot_dp = FALSE)

dat <- data.frame(round(rbind(norm$moment_fit[1:6], pois$moment_fit[1:6], 
                              nbinom$moment_fit[1:6], 
                              x$moment_fit[1:6],  norm$moment_precision[1:6]), 2))


dat <- cbind(c('Normal', 'Poisson', 'NegBin', 'ZINB', 'Moment precision'), dat)
colnames(dat) <- c('Distribution', 'mu1', 'mu2', 'mu3', 'mu4', 'mu5', 'mu6')
table1 <- xtable(dat)

write.csv(table1, file="table2.csv")

ind <- sample(c(0, 1), n, prob = c(x$params[1], 1 - x$params[1]), replace = TRUE)
Z_est <- ind*rnbinom(n, mu = x$params[2]*x$params[3]/(1-x$params[2]), size = x$params[3]) 

plot1 <- ggplot() +   
   geom_histogram(aes(x = X, fill = 'X'), alpha = .7, bins = 35)  + 
   geom_histogram(aes(x = Z_est, fill = 'Est Z'), alpha = .7, bins = 40) + 
   theme_bw() + 
   theme(legend.position = c(0.9, 0.9), 
       legend.justification =  c(0.9, 0.9), 
       panel.grid.major = element_blank(), 
       panel.grid.minor = element_blank()) + labs(x = '', title = 'Estimated vs. observed data', fill = '') + 
   scale_fill_manual(values = c('navyblue', 'orange'))

plot2 <- ggplot() + 
  geom_histogram(aes(x = Z, fill = 'True Z'), alpha = .7, bins = 30) + 
   theme_bw() + 
   theme(legend.position = c(0.9, 0.9), 
       legend.justification =  c(0.9, 0.9),
       panel.grid.major = element_blank(), 
       panel.grid.minor = element_blank()) + 
  labs(x = '', title = 'Private data', fill = '') + 
 xlim(c(-10, 20)) + 
    scale_fill_manual(values = c('navyblue', 'orange'))
  
plot <- grid.arrange(plot1, plot2, ncol=2)

ggsave(plot, file = 'plots/figure4.pdf', width = 7, height = 4)